import cv2
import numpy as np
import os
from glob import glob

# 去除噪点
def Img1(src):
    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(src, connectivity=8, ltype=None)
    img = np.zeros((src.shape[0], src.shape[1]), np.uint8)    #创建个全0的黑背景
    for i in range(1, num_labels):
        mask = labels == i             #这一步是通过labels确定区域位置，让labels信息赋给mask数组，再用mask数组做img数组的索引
        if stats[i][4] >= 2:         # 面积 筛选
            img[mask] = 255      #面积大于区域涂白留下，小于300的涂0抹去
        else:
            img[mask] = 0
           
    return img
 
def Img2(img):
    contours, hierarch = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    area = []
    for i in range(len(contours)):
        area.append(cv2.contourArea(contours[i]))   #计算轮廓所占面积
        if area[i] < 2:                   #轮廓面积，可以自己随便调
            cv2.drawContours(img,[contours[i]],0,0,-1)         #该轮廓区域填0
            continue
    return img
 
# src = cv2.imread('images/lineFrames/1.png', 0)
# cv2.imshow('input',src)
# # Img1 的效果更好
# src1 = Img1(src)
# cv2.imshow('output1', src1)
# src2 = Img2(src)
# cv2.imshow('output2', src2)
# cv2.waitKey()

def removeOutlier(filename):
    img = cv2.imread(filename, cv2.IMREAD_GRAYSCALE)

    img_rem = Img1(img)
    
    save_filename = '%s.png' % (os.path.basename(filename).split('.')[0])
    cv2.imwrite('images/lineFrames_remove/' + save_filename, img_rem)


if __name__ == "__main__":
    # if not os.path.exists('edges'):
    #     os.makedirs('edges')
    file_list = glob('images/lineFrames/*.png')
    for filename in file_list:
        removeOutlier(filename)